using namespace std;

#include <iostream>
#include <fstream>
#include <malloc.h>

#include "types.h"
#include "rsakeys.h"
#include "base64.h"
#include "cmdargs.h"

extern cmdargs args;

void rsakeys::generateKeys()
{
	//mpz_t p, q, n, phy, e, gcd, d;
	mpz_t phy, gcd;
	mpz_class p1, q1;
	gmp_randclass r(gmp_randinit_mt);
	
	mpz_init(p);
	mpz_init(q);
	mpz_init(n);
	mpz_init(phy); 
	mpz_init(e);
	mpz_init(gcd);
	mpz_init(d);
	
	int keysize = args.keysize ? args.keysize/2 : 512;
	p1 = r.get_z_bits (keysize);
	q1 = r.get_z_bits (keysize);

	mpz_set(p, p1.get_mpz_t());
	
	mpz_set(q, q1.get_mpz_t());
	
	mpz_nextprime(p, p);
	mpz_nextprime(q, q);

	mpz_mul(n, p, q);
	
	mpz_sub_ui(p, p, 1);
	mpz_sub_ui(q, q, 1);
	mpz_mul(phy, p, q);
	
	//restore p and q
	mpz_add_ui(p, p, 1);
	mpz_add_ui(q, q, 1);

	mpz_set_ui(e, E_SEED);
	
	while(mpz_cmp(phy, e) > 0) {
		mpz_gcd(gcd, phy, e);
		if (mpz_cmp_ui(gcd, 1) == 0)
			break;
		mpz_add_ui(e, e, 1);
	}

	
	if (mpz_invert(d, e, phy) == 0)
		cout << "inverse not found" << endl;
	
	
	// clear bignums.
	mpz_clear(phy);
	mpz_clear(gcd);
}


bool rsakeys::createPrivKey()
{
	generateKeys();
	uchar buff [8096];
	uint buffindex = 0;
	
	memset(buff, 0, sizeof(buff));

	//Algorithm version
	buff[buffindex++] = IDENT_INTEGER;
	buff[buffindex++] = ALGOVERLEN;
	buff[buffindex++] = ALGOVER;

	uchar temp[512];
	
	uint count = 0;
	// modulus (n)
	mpz_export (temp, &count, 1, 1, 1, 0, n);
	//cout << "# of octets in modulus: " << count << endl;
	fillIntegerBuffer(buff, buffindex, temp, count);
	
	// Public Exponent (e)
	mpz_export ((void *)temp, &count,  1, 1, 1, 0, e);
	//cout <<" # of octets in public exponent: " << count << endl;
	fillIntegerBuffer(buff, buffindex, temp, count);
	
	// Private Exponent (d)
	mpz_export (temp, &count,  1, 1, 1, 0, d);
	//cout <<" # of octets in private exponent: " << count << endl;
	fillIntegerBuffer(buff, buffindex, temp, count);
	
	// Prime1
	mpz_export ((void *)temp, &count,  1, 1, 1, 0, p);
	//cout <<" # of octets in prime1 " << count << endl;
	fillIntegerBuffer(buff, buffindex, temp, count);
	
	// Prime2
	mpz_export ((void *)temp, &count,  1, 1, 1, 0, q);
	//cout <<" # of octets in prime2: " << count << endl;
	fillIntegerBuffer(buff, buffindex, temp, count);
	
	// exponent1 d mod(p-1)
	mpz_t minus1, ex;
	mpz_init(minus1);
	mpz_sub_ui(minus1, p, 1);

	mpz_init(ex);
	mpz_mod(ex, d, minus1);

	mpz_export ((void *)temp, &count,  1, 1, 1, 0, ex);
	//cout <<" # of octets in exponent1: " << count << endl;
	fillIntegerBuffer(buff, buffindex, temp, count);
	
	// exponent2 d mod(q-1) 
	mpz_sub_ui(minus1, q, 1);
	mpz_mod(ex, d, minus1);

	mpz_export ((void *)temp, &count,  1, 1, 1, 0, ex);
	//cout <<" # of octets in exponent2: " << count << endl;
	fillIntegerBuffer(buff, buffindex, temp, count);
	
	
	// coefficient qinv mod p.
	mpz_invert(ex, q, p);
	mpz_export ((void *)temp, &count,  1, 1, 1, 0, ex);
	//cout <<" # of octets in coefficient: " << count << endl;
	fillIntegerBuffer(buff, buffindex, temp, count);
	
	mpz_clear(minus1);
	mpz_clear(ex);

	clearBignums();

	if (FALSE == writePrivateKeyToFile(buff, buffindex))
	{
		cout << "unable to write private key to file" << endl;
		return FALSE;
	}
	else
		return TRUE;

}

void rsakeys::fillIntegerBuffer(uchar *buff, uint &buffindex, uchar *temp, uint count)
{
	buff[buffindex++] = TAG_INTEGER;

//check is MSB is set;
	if (temp[0] & 0x80)
		count ++;

	if (count <= 127)
	{
		buff[buffindex++] = count; //short form
	}
	else if (count <=255) // long form
	{
			buff[buffindex++] = 0x81; // first length octet
			buff[buffindex++] = count; //number of length octets
	}
	
	if (temp[0] & 0x80)
	{
		buff[buffindex++] = '\0';
		count --;
	}

		memcpy(&buff[buffindex], temp, count);
	buffindex += count;
}


bool rsakeys::writePrivateKeyToFile(uchar *buff, uint bufflen)
{
	// construct header of private key.
	uchar hdr[10];
	uint hdrsize = 0;

	hdr[hdrsize++] = IDENT_SEQ;
	if (bufflen <=127) // one length octet
		hdr[hdrsize++] = bufflen;
	else if (bufflen <=255) // two length octets
	{
		hdr[hdrsize++] = 0x81;
		hdr[hdrsize++] = bufflen; 
	}
	else if (bufflen <= (16*16*255 + 255)) // 3 length octets
	{
		hdr[hdrsize++] = 0x82;
		hdr[hdrsize++] = (char)(bufflen/256);
		hdr[hdrsize++] = (char)(bufflen % 256);
	}

	string key((char *)hdr, hdrsize);
	key += string ((char *)buff, bufflen);
	
	base64 b;
	string key2 = b.encodeBase64(key);

	if (args.outfile.empty())
	{
		// print key to console
		cout << priv_key_start2 << endl;
		for (uint i = 0; i < key2.size(); i += 64)
			 cout << key2.substr(i, 64).c_str() << endl;
		cout << priv_key_end2 << endl;
	}
	else
	{
		// save key to file
		ofstream os(args.outfile.c_str(), ios::binary | ios::app);
		if (!os.is_open())
		{
			cout << "Error: cannot open output file: " << args.outfile << endl;
			return FALSE;
		}
		os << priv_key_start2 <<endl;

		for (uint i = 0; i < key2.size(); i += 64)
			 os << key2.substr(i, 64).c_str() << endl;

		os << priv_key_end2 << endl;
		os.close();

	}
	
	return TRUE;
}

void rsakeys::clearBignums()
{
	mpz_clear(p);
	mpz_clear(q);
	mpz_clear(n);
	mpz_clear(e);
}

bool rsakeys::createPubKey()
{
	cout << "creating public key" << endl;
	base64 b;

	ifstream is(args.infile.c_str(), ios::binary);
	if (!is.is_open())
	{
		cout << "Error: cannot open input private key file: " << args.infile << endl;
		return FALSE;
	}

	ofstream os (args.outfile.c_str(), ios::binary | ios::app);
	if (!os.is_open())
	{
		cout << "ERROR: cannot create public key file: " << args.outfile << endl;
		return FALSE;
	}
	
	// If input is in PEM, convert to DER. 
	if (args.infile.find(".der") == string::npos )
	{
		string derfile = args.infile;
		derfile.replace(derfile.end()-3, derfile.end(), "der");
		
		if (!b.decodeKeyFromBase64(args.infile.c_str(), derfile.c_str()))
		{
			cout << "unable to convert PEM to DER " << endl;
			return FALSE;
		}
		args.infile = derfile;
		is.close();
		is.open(args.infile.c_str(), ios::binary);
		if (!is.is_open())
		{
			cout << "ERROR openning DER converted private key file" << endl;
			return FALSE;
		}
	}

	uchar *inbuff; // to read from the private file

	// read the complete private key file
	is.seekg(0, ios::end);
	uint len = is.tellg();
	is.seekg(0, ios::beg);
	inbuff = (uchar *)malloc(len); // buffer to read in the private key file.
	is.read((char *)inbuff, len);
	is.close();
	
	uint index = 0;
	uint noflenoctets = 0;
	uint contentoctets = 0;
	
	// skip header.
	index = 1;
	uchar firstlenoctet = inbuff[index++];
	if (firstlenoctet & MULTLENOCTET)
	{
		noflenoctets = firstlenoctet & 0x7f;
		index += noflenoctets;
	}
	
	// skip the version info.
	index += 3;

	//reached modulus.
	if (inbuff[index] != 2)
	{
		cout<< "error in private file" << endl;
		return FALSE;
	}
	uint readindex = index;
	uint modpluse = 0;
	firstlenoctet = inbuff[index + 1];
	if (firstlenoctet & MULTLENOCTET)
	{
		noflenoctets = firstlenoctet & 0x7f;
		if (noflenoctets == 1)
			contentoctets = inbuff[index+2];
		else if (noflenoctets == 2)
			contentoctets = inbuff[index+2]*256 + inbuff[index+3];
			
		noflenoctets += 1;
	}
	else
	{
		noflenoctets = 1;
		contentoctets = firstlenoctet;
	}
	modpluse = 1 + noflenoctets + contentoctets;

	// public component.
	index += modpluse;

	firstlenoctet = inbuff[index + 1];
	if (firstlenoctet & MULTLENOCTET)
	{
		noflenoctets = firstlenoctet & 0x7f;
		if (noflenoctets == 1)
			contentoctets = inbuff[index+2];
		else if (noflenoctets == 2)
			contentoctets = inbuff[index+2]*256 + inbuff[index+3];
			
		noflenoctets += 1;
	}
	else
	{
		noflenoctets = 1;
		contentoctets = firstlenoctet;
	}
	modpluse += 1 + noflenoctets + contentoctets;

// create public key buffer.

	char  seqheader1[5];
	int seqheader1len = 0;
	seqheader1[seqheader1len++] = IDENT_SEQ;
	if ( modpluse <= 127) // one length octet
		seqheader1[seqheader1len++] = modpluse;

	else if (modpluse <= 255) // two length octets
	{
		seqheader1[seqheader1len++] = 0x81;
		seqheader1[seqheader1len++] = modpluse;
	}
	else if (modpluse <= (16*16*255 + 255)) // 3 length octets
	{
		seqheader1[seqheader1len++] = 0x82;
		seqheader1[seqheader1len++] = (char)(modpluse / 256);
		seqheader1[seqheader1len++] = (char)(modpluse % 256);
	}

	char bitstring[5];
	int bitstrlen = 0;
	int lentoencode = modpluse + seqheader1len + 1;// +1 for the pad byte.

	bitstring[bitstrlen++] = IDENT_BITSTRING_PRIM;
	if ( lentoencode <= 127) // one length octet
		bitstring[bitstrlen++] = lentoencode;

	else if ( lentoencode <= 255) // two length octets
	{
		bitstring[bitstrlen++] = 0x81;
		bitstring[bitstrlen++]= lentoencode;
	}
	else if (lentoencode <= (16*16*255 + 255)) // 3 length octets
	{
		bitstring[bitstrlen++] = 0x82;
		bitstring[bitstrlen++] = (char)(lentoencode / 256);
		bitstring[bitstrlen++] = (char)(lentoencode % 256);
	}
	bitstring[bitstrlen++] = '\0'; //padding
	
	char pubkeyinascii[512];
	index = 0;

	int seqlength = sizeof(pubkeyheader) + bitstrlen + seqheader1len + modpluse;

// write initial sequence.
	pubkeyinascii[index++] = IDENT_SEQ;
	if (seqlength <=127) // one length octet
		pubkeyinascii[index++] = seqlength;
	else if (seqlength <= 255) // two length octets
	{
		pubkeyinascii[index++] = 0x81;
		pubkeyinascii[index++] = seqlength; 
	}
	else if (seqlength <= (16*16*255 + 255)) // 3 length octets
	{
		pubkeyinascii[index++] = 0x82;
		pubkeyinascii[index++] = (char)(seqlength / 256);
		pubkeyinascii[index++] = (char)(seqlength % 256);
	}
	//copy pub key object identifier
	memcpy (&pubkeyinascii[index], pubkeyheader, sizeof(pubkeyheader));
	index += sizeof(pubkeyheader);

	//copy the bitstring.
	memcpy (&pubkeyinascii[index], bitstring, bitstrlen);
	index += bitstrlen;
	
	// copy the sequence
	memcpy (&pubkeyinascii[index], seqheader1, seqheader1len);
	index += seqheader1len;

	//copy modulus and pub key.
	memcpy (&pubkeyinascii[index], &inbuff[readindex], modpluse);
	index += modpluse;

	string pubkeyinpem;
	pubkeyinpem = b.encodeBase64(string (pubkeyinascii, index));

	// write to pub key file.
	os << pub_key_start << endl;

	for (uint i = 0; i < pubkeyinpem.size(); i += 64)
		 os << pubkeyinpem.substr(i, 64).c_str() << endl;

	os << pub_key_end << endl;

	os.close();
	free (inbuff);
	return TRUE;

}

bool rsakeys::convertPrivKeyFile()
{
	base64 b;

	if (!args.inform.compare("PEM") && !args.outform.compare("DER"))
	{
		// convert PEM to DER
		if (b.decodeKeyFromBase64(args.infile.c_str(), args.outfile.c_str()))
			return TRUE;
	}
	else if (!args.inform.compare("DER") && !args.outform.compare("PEM"))
	{
		// convert DER to PEM
		if (b.encodeKeyToBase64(args.infile.c_str(), args.outfile.c_str()))
			return TRUE;
	}

	return FALSE;
}

